package com.hzxgs.tracelogstarter.trace;

import com.hzxgs.tracelogstarter.properties.TraceProperties;
import com.hzxgs.tracelogstarter.util.TraceUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.amqp.AmqpException;
import org.springframework.amqp.core.Correlation;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.core.MessagePostProcessor;
import org.springframework.amqp.core.MessageProperties;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.Objects;

@Component
@ConditionalOnClass(value = { MessagePostProcessor.class })
@Slf4j
public class RabbitMQAfterReceiveMessagePostProcessor implements MessagePostProcessor {

    @Resource
    private TraceProperties traceProperties;

    @Override
    public Message postProcessMessage(Message message) throws AmqpException {
        this.putTraceFromMessage(message);
        return message;
    }

    @Override
    public Message postProcessMessage(Message message, Correlation correlation) {
        return MessagePostProcessor.super.postProcessMessage(message, correlation);
    }

    @Override
    public Message postProcessMessage(Message message, Correlation correlation, String exchange, String routingKey) {
        if (this.putTraceFromMessage(message)) {
            return message;
        } else {
            return MessagePostProcessor.super.postProcessMessage(message, correlation, exchange, routingKey);
        }
    }

    private Boolean putTraceFromMessage(Message message) {
        if (!traceProperties.getEnable()) {
            return false;
        }
        try {
            MessageProperties messageProperties = message.getMessageProperties();
            Object traceIdHeader = messageProperties.getHeader(TraceUtils.TRACE_ID_HEADER);
            Object spanIdHeader = messageProperties.getHeader(TraceUtils.SPAN_ID_HEADER);
            String traceId = "";
            String spanId = "";
            if (!Objects.isNull(traceIdHeader)) {
                traceId = String.valueOf(traceIdHeader);
            }
            if (!Objects.isNull(spanIdHeader)) {
                spanId = String.valueOf(spanIdHeader);
            }
            // 设置MDC
            TraceUtils.putTrace(traceId, spanId);

            return true;
        } catch (Exception e) {
            log.error("rabbitMq设置traceId失败: {}", e.getMessage());
            e.printStackTrace();
            return false;
        }
    }
}
